import cv2
import numpy as np
import pygetwindow as gw
import pyautogui
from ultralytics import YOLO

# 加载 YOLO 模型并指定使用 GPU
model = YOLO(r"H:\yol\yolo\v11\runs\detect\train42\weights\best.pt")
model.to('cuda')

# 查找标题包含“梦幻西游 ONLINE”的窗口
windows = gw.getWindowsWithTitle('梦幻西游 ONLINE')

# 定义检测窗口的尺寸
target_width = 1024
target_height = 768

for window in windows:
    if window.isMinimized:
        window.restore()  # 恢复最小化的窗口
    window.activate()  # 激活窗口

    left, top, width, height = window.left, window.top, window.width, window.height

    while True:
        # 截取窗口区域的屏幕截图
        screenshot = pyautogui.screenshot(region=(left, top, width, height))
        screenshot = np.array(screenshot)
        screenshot = cv2.cvtColor(screenshot, cv2.COLOR_RGB2BGR)

        # 调整图像尺寸到目标大小
        resized_screenshot = cv2.resize(screenshot, (target_width, target_height))

        # 进行目标检测
        results = model(resized_screenshot)

        # 绘制检测框
        annotated_frame = results[0].plot()

        # 显示结果
        cv2.imshow(f'Detection Result - {window.title}', annotated_frame)

        # 按 'q' 键退出循环
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cv2.destroyAllWindows()
